/*
 *   Copyright (C) 2009 Ultratronik Entwicklungs GmbH 
 *
 *     - Author: Sören Wellhöfer <soeren.wellhoefer@gmx.net>
 *
 *   This program is free software; you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation; either version 2 of the License, or
 *   (at your option) any later version.
 *
 *
 * Compilation:
 *  gcc -o cpuload cpuload.c -lm -lcurses
 *
 * Notes:
 *   Run as root.
 *
 * Abstract:
 *   Attempts to impose a steady cpu load on the system by dynamically
 *   scaling a hot-loop depending on the current load or utilization. 
 *   The algorithm used narrows the variance down to about
 *   0.05. If cpu utilization rather than load is controlled consider
 *   using `-method exp' as it does not seem to nutate as much.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <unistd.h>
#include <fcntl.h>
#include <math.h>
#include <errno.h>
#include <stropts.h>
#include <poll.h>
#include <curses.h>
#include <term.h>
#include <sys/time.h>
#include <sys/resource.h>

const char *usg = "Steady CPU load generator\n"
"usage: %s [-verbose] [-poll MSEC] [-nice VAL] [-method exp|sqrt] [-scale VAL] [-util] LOAD%\n\n";
const char *prog;

/* Constants */
#define DEFAULT_POLL_INTERVAL 500 * 1000 /* .5s */
#define DEFAULT_SCALE_FACT 10 * 100 /* .1s */

/* Macros */
#define sign(x) (((x) > 0) - ((x) < 0))
#define abs(x) ((x) < 0 ? -(x) : (x))

/* Gobals */
float (*cpu_current)();
int poll_interval, flags;
float pivot_load;
pid_t child;
int fdch;
int pnice;
int scale;

/* Prototypes */
static void* memalloc(size_t);
static void cpuload();
static void balance();

/* Definitions */
struct {
	void (*f)();
	const char *m;
} oop[] = {
	{&cpuload, "r"},
	{&balance, "w"}
};

enum {
	READFD = 0,
	WRITEFD,

	FLAG_VERBOSE = 1<<0,
	FLAG_M_EXP = 1<<1,
	FLAG_M_SQRT = 1<<2,
	FLAG_TERMCAP = 1<<3,
};

/* Implementation */
static inline void usage()
{
	printf(usg, prog);
}

static void errmsg(char *frm, int with_errno, ...)
{
	static const char *err = "error";
	static const char *format = "%s: %s\n";
	size_t len;
	char *pre;
	va_list l;

	va_start(l, with_errno);

	len = strlen(err) + strlen(format) + strlen(frm);
	pre = (char*)memalloc((len + 1) * sizeof(char));
	pre[len] = 0;
	
	sprintf(pre, format, err, frm);
	vfprintf(stderr, pre, l);

	if(with_errno) fprintf(stderr, "(%s)\n", strerror(errno));

	free(pre);
	va_end(l);
}

static void* memalloc(size_t s)
{
	void *mem = (void*)malloc(s);
	if(!mem) {
		fprintf(stderr, "out of memory\n");
		exit(-1);
	}
	return mem;
}

static float load_current()
{
	static const char *lavg = "/proc/loadavg";
	float laverage;	
	FILE *f;
	
	if(!(f = fopen(lavg, "r"))) {
		errmsg("could not open %s", 1, lavg);
		exit(-1);
	}

	if(fscanf(f, "%f", &laverage) != 1) {
		errmsg("could not obtain load average", 1);
		exit(-1);
	}

	fclose(f);
	return laverage;
}

static float util_current()
{
	static const char *cutil = "/proc/stat";
	char buf[128], *p;
	int n, k, val;
	FILE *f;

	if(!(f = fopen(cutil, "r"))) {
		errmsg("could not open %s", 1, cutil);
		exit(-1);
	}
	
	int total_time, idle_time, idle_diff, total_diff;
	static int idle_prev, total_prev;
	
	fgets((p = buf), sizeof(buf), f);
	fscanf(f, "%*s%n", &n), p += n;
	
	for(total_time = k = 0; sscanf(p, "%d%n", &val, &n); k++, p += n) {
		total_time += val;
		if(k == 3) idle_time = val;
	}

	idle_diff = idle_time - idle_prev;
	total_diff = total_time - total_prev;
	idle_prev = idle_time;
	total_prev = total_time;

	fclose(f);
	return (float)(total_diff - idle_diff) / (float)total_diff;
}

/* The work-horse; a controllable hot-loop */
static void cpuload()
{
	int slval = 0, inc;
	struct pollfd fds[1];

	fds[0].fd = fdch;
	fds[0].events = POLLIN;
	
	for(;;) {
		if(poll(fds, 1, 0)) {
			if(read(fdch, &inc, sizeof(inc)) < 0) {
				errmsg("could not read pipe", 1);
				_exit(-1);
			}
		} else
			inc = 0;
		
		slval += inc;
		if(slval <= 0) slval = 0;
		else usleep(slval);
	}
}

static void balance()
{
	float current, diff;
	int inc, b;
	
	if((setpriority(PRIO_PROCESS, child, pnice) == -1)) {
		errmsg("cannot set process priority", 1);
		exit(-1);
	}
	
	for(b = 0;; b++) {
		current = cpu_current();
		diff = current - pivot_load;
		inc = sign(diff) * (flags & FLAG_M_SQRT ? sqrt(abs(diff)) : diff * diff) * scale;
		
		if(flags & FLAG_VERBOSE) {
			if(flags & FLAG_TERMCAP) tputs(tparm(tigetstr("hpa"), 0), 1, &putchar);
			fprintf(stdout, "load=%.2f, diff=%.2f, inc=%d  (%d)  %c",
				current, diff, inc, b, flags & FLAG_TERMCAP ? ' ' : '\n');
		}
		
		if(write(fdch, &inc, sizeof(inc)) < 0) {
			errmsg("pipe channel deceased", 1);
			exit(-1);
		}
		
		usleep(poll_interval);
	}
}

static void start()
{
	int fdx, fds[2];

	if(pipe(&fds[0]) < 0) {
		errmsg("cannot create pipe", 1);
		exit(-1);
	}

	if((child = fork()) < 0) {
		errmsg("cannot fork", 1);
		exit(-1);
	}

	fdx = child ? WRITEFD : READFD;
	close(fds[fdx ^ 1]);
	fdch = fds[fdx];
	oop[fdx].f();
}

int main(int argc, char **argv)
{
	int terr;
	
	prog = *argv++;
	
	if(argc < 2) {
		usage();
		return -1;
	}

	flags = pnice = 0;
	cpu_current = &load_current;
	poll_interval = DEFAULT_POLL_INTERVAL;
	scale = DEFAULT_SCALE_FACT;

	setvbuf(stdout, NULL, _IONBF, 0);
	
	if(setupterm((char*)0, fileno(stdout), &terr) == OK)
		flags |= FLAG_TERMCAP;
	else
		errmsg("terminal capabilities not initialized (code %d)", 0, terr);
	
	for( ; --argc > 0; argv++) {

		if(!strcmp("-h", *argv) || !strcmp("-help", *argv)) {
			usage();
			return 0;
		} else if(!strcmp("-poll", *argv)) {
			poll_interval = atoi(*++argv) * 1000;
			--argc;
		} else if(!strcmp("-scale", *argv)) {
			scale = atoi(*++argv);
			--argc;
		} else if(!strcmp("-verbose", *argv)) {
			flags |= FLAG_VERBOSE;
		} else if(!strcmp("-util", *argv)) {
			cpu_current = &util_current;
		} else if(!strcmp("-nice", *argv)) {
			pnice = atoi(*++argv);
			--argc;
		} else if(!strcmp("-method", *argv)) {
			if(!strncmp(*++argv, "exp", strlen(*argv)))
				flags |= FLAG_M_EXP;
			else if(!strncmp(*argv, "sqrt", strlen(*argv)))
				flags |= FLAG_M_SQRT;
			else {
				errmsg("unknown scaling method %s", 0, *argv);
				exit(-1);
			}
			--argc;
		} else {
			pivot_load = (float)atof(*argv);
		}
	}

	if(pivot_load <= 0) {
		errmsg("load must be > 0", 0, *argv);
		return -1;
	}

	start();
	close(fdch);

	return 0;
}
